import cv2
import numpy as np
import os
from sklearn.cluster import KMeans

class img_features(object):

    def __init__(self):
        pass

    def get_resize_img(self, src):
        img = cv2.imread(src)
        return self.resize_img(img)
    
    def get_gray_hist_img(self, img):
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return cv2.equalizeHist(img_gray)

    def get_gray_img(self, img):
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    def get_hist_img(self, img):
        yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
        y = yuv[:,:,0]
        yuv[:,:,0] = cv2.equalizeHist(y)
        bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
        return bgr

    def resize_img(self, img):
        return cv2.resize(img, (1024, 768))
        
    def show_img(self, desc, img):
        cv2.imshow(desc, img)

    def get_star_points(self, img):
        star = cv2.xfeatures2d.StarDetector_create()
        return star.detect(img)

    def get_sift_features(self, img, points):
        sift = cv2.xfeatures2d.SIFT_create()
        points, features = sift.compute(img, points)
        return points, features

    def get_features(self, gray_img):
        points = self.get_star_points(gray_img)
        return self.get_sift_features(gray_img, points)
    
    def draw_keypoints(self, img, dest_img, points):
        cv2.drawKeypoints(img, points, dest_img, (0,0,255), flags=cv2.DRAW_MATCHES_FLAGS_DEFAULT)

    def img_write(self, img_name, img):
        cv2.imwrite("./result/" + img_name, img)

    def close_all_windows(self):
        key = cv2.waitKey(0)
        if('q' == ord(key)):
            cv2.destroyAllWindows()

    def corner_harris(self, img):
        return cv2.cornerHarris(img, 10, 5, 0.01)

def classify():
    path_dir = os.listdir("./all_data")
    all_files = []
    x_data = []
    img_f = img_features()
    for f in path_dir:
        print("reading file " + f)
        org_img = img_f.get_resize_img("./all_data/" + f)
        gray_img = img_f.get_gray_img(org_img)
        points, features = img_f.get_features(gray_img)
        all_files.append(f)
        reshape = features.reshape(-1,1)
        x_data.append(reshape[0].tolist())
        x = np.array(x_data)
    
    kms = KMeans(n_clusters = 2)
    y = kms.fit_predict(x)
    result = {}
    for i in range(len(all_files)):
        result[all_files[i]] = str(y[i])
        print(all_files[i] + " --> " + str(y[i]))
    

def test():
    img_path = "./data/dog/1.jpg"
    img_f = img_features()

    org_img = img_f.get_resize_img(img_path)
    #img_f.show_img("org img", org_img)
    img_f.img_write("org_img.jpg", org_img)

    gray_img = img_f.get_gray_img(org_img)
    #img_f.show_img("gray img", gray_img)
    img_f.img_write("gray_img.jpg", gray_img)

    gray_hist_img = img_f.get_gray_hist_img(org_img)
    #img_f.show_img("gray hist img", gray_hist_img)
    img_f.img_write("gray_hist_img.jpg", gray_hist_img)

    org_points = img_f.get_star_points(org_img)
    gray_points = img_f.get_star_points(gray_img)

    points, features = img_f.get_sift_features(gray_img, gray_points)
    gray_hist_points, gray_hist_features = img_f.get_features(gray_hist_img)

    org_img1 = img_f.get_resize_img(img_path)
    img_f.draw_keypoints(org_img, org_img, org_points)
    img_f.draw_keypoints(org_img1, org_img1, gray_hist_points)

    #img_f.show_img("org points", org_img)
    img_f.img_write("org_img_star_points.jpg", org_img)
    #img_f.show_img("gray hist points", org_img1)
    img_f.img_write("org_img_more_points.jpg", org_img1)

    corner_img = img_f.corner_harris(gray_img)
    org_img[corner_img>0.03*corner_img.max()] = [0,0,0]

    #img_f.show_img("corner img", org_img)
    img_f.img_write("corner_img.jpg", org_img)

test()
classify()
